package cn.swing.main.srv.cv.utils;

import cn.swing.main.common.model.Tuple2;
import cn.swing.main.common.model.Tuple3;
import cn.swing.main.srv.cv.model.TennisPoint;
import cn.swing.main.srv.cv.model.Trajectory;
import cn.swing.main.srv.cv.tracker.AdvancedTracker;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.springframework.util.CollectionUtils;

/**
 * 合并这些轨迹
 * 步骤1 根据overlap和gap判断当前轨迹的候选轨迹，overlap不超过15的或者gap不超过20的，都认为是候选轨迹
 * 步骤2 根据候选轨迹和当前轨迹的匹配度判断是否连接
 * 判断规则（要考虑到vx vy速度的变化率）：
 * 1. 重合轨迹：取t1的最后5点，与t2的非重合的前5点进行匹配，匹配度大于0.8，则认为连接。如果不行，则取t2的前5点，与t1的非重合的最后5点进行匹配，匹配度大于0.8，则认为连接。
 * 2. 断点轨迹：取t1的最后5点，与t2的前5点进行匹配，匹配度大于0.8，则认为连接。如果不行，则取t2的前5点，与t1的最后5点进行匹配，匹配度大于0.8，则认为连接。
 */
public class TrajectoryMerger {

    /**
     * 合并轨迹时，两段轨迹帧差不能超过{x}
     */
    public static final double MERGE_TRAJ_FRAME_GAP_THRESHOLD = 30;

    /**
     * 合并轨迹时，两段轨迹重叠不能超过{x}
     */
    public static final double MERGE_TRAJ_FRAME_OVERLAP_THRESHOLD = 20;

    // 合并主逻辑
    public List<Trajectory> mergeTrajectories(List<Trajectory> input) {
        List<Trajectory> result = new ArrayList<>();
        input.sort(Comparator.comparingInt(Trajectory::getStartFrameNumber));

        // 存放被连起来的轨迹 key为end部分，value为连接的轨迹，以及重合类型
        Map<Trajectory, Tuple3<Trajectory, Trajectory, Integer>> merged = new HashMap<>();
        for (int i = 0; i < input.size(); i++) {
            Trajectory current = input.get(i);

            // 寻找候选轨迹
            List<Trajectory> candidates = findCandidates(current, input.subList(i + 1, input.size()));
            if (candidates.isEmpty()) {
                continue;
            }
            // 类型：1：重合轨迹，保留前段最后5个 舍弃后段重复的轨迹
            // 2：重合轨迹，保留后段前面5个 舍弃前段重复的轨迹
            // 3：断点不重合
            int[] type = new int[candidates.size()];
            // 记录每个候选轨迹的匹配度 值越小越匹配
            double[] distance = new double[candidates.size()];
            for (int c = 0; c < candidates.size(); c++) {
                Trajectory candidate = candidates.get(c);

                // 判断重叠类型
                boolean isOverlap = current.getEndFrameNumber() >= candidate.getStartFrameNumber();

                if (isOverlap) {
                    // 重合轨迹匹配
                    // 获取轨迹1的最后5个点，与轨迹2的非重合的前5个点进行匹配
                    List<TennisPoint> t1Points = current.last5();
                    List<TennisPoint> t2Points = candidate.getFirstNonOverlapPoints(current, 5);
                    if (t1Points.isEmpty() || t2Points.isEmpty()) {
                        break;
                    }
                    double zx = calculateMatchScore(t1Points, t2Points, false);
                    // t1Points 和 t2Points平均X Y 之间的曼哈顿距离
                    final double zxm = ManhattanDistance.calculateManhattanDistance(
                            t1Points.stream().mapToDouble(TennisPoint::getX).average().getAsDouble(),
                            t1Points.stream().mapToDouble(TennisPoint::getY).average().getAsDouble(),
                            t2Points.stream().mapToDouble(TennisPoint::getX).average().getAsDouble(),
                            t2Points.stream().mapToDouble(TennisPoint::getY).average().getAsDouble()
                    );

                    // 获取轨迹2的前5个点，与轨迹1的非重合的最后5个点进行匹配
                    t1Points = current.getLastNonOverlapPoints(candidate, 5);
                    t2Points = candidate.first5();
                    if (t1Points.isEmpty() || t2Points.isEmpty()) {
                        break;
                    }
                    double fx = calculateMatchScore(t1Points, t2Points, false);
                    // t1Points 和 t2Points平均X Y 之间的曼哈顿距离
                    final double fxm = ManhattanDistance.calculateManhattanDistance(
                            t1Points.stream().mapToDouble(TennisPoint::getX).average().getAsDouble(),
                            t1Points.stream().mapToDouble(TennisPoint::getY).average().getAsDouble(),
                            t2Points.stream().mapToDouble(TennisPoint::getX).average().getAsDouble(),
                            t2Points.stream().mapToDouble(TennisPoint::getY).average().getAsDouble()
                    );

                    if (Math.min(zx, zxm) > Math.min(fx, fxm)) {
                        distance[c] = Math.min(fx, fxm);
                        type[c] = 1;
                    } else {
                        distance[c] = Math.min(zx, zxm);
                        type[c] = 2;
                    }
                } else { // 断点轨迹匹配
                    List<TennisPoint> t1Points = current.getPoints().subList(
                            Math.max(0, current.getPoints().size() - 5), current.getPoints().size());
                    List<TennisPoint> t2Points = candidate.getPoints().subList(0, Math.min(5,
                            candidate.getPoints().size()));
                    if (t1Points.isEmpty() || t2Points.isEmpty()) {
                        break;
                    }
                    // 断点不涉及去重，正反能匹配上就行，类型都是3
                    double zx = calculateMatchScore(t1Points, t2Points, false);
                    double fx = calculateMatchScore(t1Points, t2Points, true);
                    // t1Points 和 t2Points平均X Y 之间的曼哈顿距离 / 2 代表单边
                    final double m = ManhattanDistance.calculateManhattanDistance(
                            t1Points.stream().mapToDouble(TennisPoint::getX).average().getAsDouble(),
                            t1Points.stream().mapToDouble(TennisPoint::getY).average().getAsDouble(),
                            t2Points.stream().mapToDouble(TennisPoint::getX).average().getAsDouble(),
                            t2Points.stream().mapToDouble(TennisPoint::getY).average().getAsDouble()) / 2;
                    distance[c] = Math.min(m, Math.min(zx, fx));
                    type[c] = 3;
                }
            }
            // 计算score最小的值的位置
            int minIndex = 0;
            if (distance.length > 0) {
                for (int j = 1; j < distance.length; j++) {
                    if (distance[j] < distance[minIndex]) {
                        minIndex = j;
                    }
                }
                final double v = distance[minIndex];
                if (v < AdvancedTracker.MAX_DISTANCE_PER_FRAME * 2) {
                    int tp = type[minIndex];
                    // 合并轨迹
                    merged.put(candidates.get(minIndex), Tuple3.of(current, candidates.get(minIndex), tp));
                }
            }
        }
        // 合并处理
        for (int i = input.size() - 1; i >= 0; i--) {
            Trajectory trajectory = input.get(i);
            if (merged.containsKey(trajectory)) {
                // 说明这个被连接了，需要合并
                Tuple3<Trajectory, Trajectory, Integer> tuple3 = merged.get(trajectory);
                final Trajectory head = tuple3.getT1();
                final Trajectory tail = tuple3.getT2();
                switch (tuple3.getT3()) {
                    case 1:
                        // 去掉后段与前段重合的点然后拼接合并
                        tail.getPoints().removeIf(p -> p.getFrameNumber() >= head.getEndFrameNumber());
                        head.merge(tail);
                        System.out.println("去掉后段与前段重合的点然后拼接合并: " + head.getId() + " - " + tail.getId());
                        break;
                    case 2:
                        // 去掉前段与后段重合的点然后拼接合并
                        head.getPoints().removeIf(p -> p.getFrameNumber() >= tail.getStartFrameNumber());
                        head.merge(tail);
                        System.out.println("去掉前段与后段重合的点然后拼接合并: " + head.getId() + " - " + tail.getId());
                        break;
                    case 3:
                    default:
                        // 直接拼接
                        head.merge(tail);
                        System.out.println("直接拼接: " + head.getId() + " - " + tail.getId());
                        break;
                }
            } else {
                // 没有被连接，直接加入结果集
                result.add(trajectory);
            }
        }

        return result;
    }

    private List<Trajectory> findCandidates(Trajectory current, List<Trajectory> trajectories) {
        List<Trajectory> candidates = new ArrayList<>();
        for (Trajectory t : trajectories) {
            if (t == current) {
                continue;
            }

            int overlap = current.getEndFrameNumber() - t.getStartFrameNumber() + 1;
            int gap = t.getStartFrameNumber() - current.getEndFrameNumber() - 1;

            if ((overlap >= 0 && overlap <= MERGE_TRAJ_FRAME_OVERLAP_THRESHOLD)
                    || (gap >= 0 && gap <= MERGE_TRAJ_FRAME_GAP_THRESHOLD)) {
                candidates.add(t);
            }
        }
        return candidates;
    }

    // 步骤2：计算匹配度
    private double calculateMatchScore(List<TennisPoint> points1, List<TennisPoint> points2, boolean isReverse) {
        if (points1.size() < 3 || points2.size() < 3) {
            return Integer.MAX_VALUE;
        }

        if (!isReverse) {
            // 正向处理 获取第二段的第一个 p0->p1->p2
            TennisPoint p2 = points2.get(0);
            // 速度变化率计算
            final TennisPoint p1 = points1.get(points1.size() - 1);
            final TennisPoint p0 = points1.get(points1.size() - 2);
            double frameDiff = p2.getFrameNumber() - p1.getFrameNumber();
            double preX = (p1.getVx() * 2 - p0.getVx()) * frameDiff + p1.getX();
            double preY = (p1.getVy() * 2 - p0.getVy()) * frameDiff + p1.getY();

            return Math.sqrt(Math.pow(preX - p2.getX(), 2) + Math.pow(preY - p2.getY(), 2));
        } else {
            // 反向处理  p0->p1->p2
            final TennisPoint p2 = points1.get(points1.size() - 1);
            final TennisPoint p1 = points2.get(1);
            final TennisPoint p0 = points2.get(2);
            double frameDiff = p2.getFrameNumber() - p1.getFrameNumber();
            double preX = p1.getVx() * 2 - p0.getVx() + p1.getX();
            double preY = (p1.getVy() * 2 - p0.getVy()) * frameDiff + p1.getY();
            return Math.sqrt(Math.pow(preX - p2.getX(), 2) + Math.pow(preY - p2.getY(), 2));
        }
    }

}
